package com.yy.utils;

import org.apache.zookeeper.*;

import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

/**
 * @Author: chen
 * @Date: 2022/9/23 17:38
 * @Desc: 基于zookeeper实现的 自旋锁 阻塞锁 可重入锁
 * 问题: 只有锁的持有者才需要重入,持有锁的前置节点一定为空,即便重入前置节点依然是空,为什么还需要编程判断
 * 个人理解: 逻辑不受影响,性能会有提升
 */
public class ZkDistributeLock {
    ZooKeeper zooKeeper;
    String lockName;
    static final String ROOT_PATH = "/locks";
    String currentNodePath;
    static final ThreadLocal<Integer> THREAD_LOCAL = new ThreadLocal<Integer>();

    public ZkDistributeLock(ZooKeeper zooKeeper, String lockName) {
        this.zooKeeper = zooKeeper;
        this.lockName = lockName;
        // 如果根节点不存在,则创建
        // curator源码中说exists不管存不存在都会进行监听,有内存泄露的风险
        try {
            if (null == zooKeeper.exists(ROOT_PATH, false)) {
                zooKeeper.create(ROOT_PATH, null, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
            }
        } catch (KeeperException | InterruptedException e) {
            e.printStackTrace();
        }
    }

    /**
     * 如果加锁失败,自旋等待
     */
    @Deprecated
    public void lock() {
        while (!tryLock()) {
            try {
                TimeUnit.MILLISECONDS.sleep(100);
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * 加锁方法,创建临时节点,为了防止服务器宕机带来的死锁问题
     */
    private boolean tryLock() {
        try {
            zooKeeper.create(ROOT_PATH + "/" + lockName, null, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL);
            return true;
        } catch (KeeperException | InterruptedException e) {
            return false;
        }
    }

    /**
     * 解锁方法,删除节点
     */
    @Deprecated
    public void unlock() {
        try {
            zooKeeper.delete(ROOT_PATH + "/" + lockName, -1);
        } catch (InterruptedException | KeeperException e) {
            e.printStackTrace();
        }
    }

    /**
     * 阻塞的公平锁,非自旋,自旋会有性能开销
     */
    public void fairLock() {
        try {
            //判断threadLocal中是否有锁,有锁直接进行重入
            Integer reentryNum = THREAD_LOCAL.get();
            if (null != reentryNum && reentryNum > 0) {
                THREAD_LOCAL.set(reentryNum + 1);
                return;
            }

            currentNodePath = zooKeeper.create(ROOT_PATH + "/" + lockName + "-", null, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL_SEQUENTIAL);
            // 获取强制节点,如果前置节点为空,则获取锁成功,否则监听前置节点
            String preNode = getPreNode();

            if (null != preNode) {
                // 利用闭锁思想,实现阻塞功能
                CountDownLatch countDownLatch = new CountDownLatch(1);
                // 再次判断zk中的前置节点是否存在,因为获取前置节点的操作不具备原子性,如果不存在直接获取锁
                if (null == zooKeeper.exists(ROOT_PATH + "/" + preNode, new Watcher() {
                    @Override
                    public void process(WatchedEvent watchedEvent) {
                        countDownLatch.countDown();
                    }
                })) {
                    THREAD_LOCAL.set(1);
                    return;
                }

                countDownLatch.await();
            }

            THREAD_LOCAL.set(1);
        } catch (KeeperException | InterruptedException e) {
            throw new IllegalMonitorStateException("create sequential node fail to lock");
        }
    }

    /**
     * 解锁方法,删除节点
     */
    public void fairUnlock() {
        try {
            int postDeductNum = THREAD_LOCAL.get() - 1;
            THREAD_LOCAL.set(postDeductNum);
            if (postDeductNum == 0) {
                zooKeeper.delete(currentNodePath, -1);
            }
        } catch (InterruptedException | KeeperException e) {
            e.printStackTrace();
        }
    }

    /**
     * 获取当前节点的前一个节点
     * 为什么不直接操作字符串,反而获取节点列表
     * 因为根节点下面可能有很多锁,比如
     * lockx-123  locky-123  lockz-123
     * 并不能保证每个锁是连续递增的
     *
     * @return 上一个节点路径
     */
    private String getPreNode() {
        try {
            String currentNodeName = currentNodePath.substring(currentNodePath.lastIndexOf("/") + 1);

            List<String> sortedChildrenNodes = zooKeeper.getChildren(ROOT_PATH, false).stream()
                    .filter(node -> node.startsWith(lockName + "-"))
                    .sorted()
                    .collect(Collectors.toList());

            // 获取当前节点在列表中的下标
            int index = sortedChildrenNodes.indexOf(currentNodeName);

            // 如果在列表中没有找到当前节点或者列表为空,则抛出异常
            if (index < 0 || sortedChildrenNodes.isEmpty()) {
                throw new IllegalMonitorStateException(
                        String.format("can not find current node %s in %s",
                                currentNodeName,
                                sortedChildrenNodes));
            }

            // 如果当前节点不是第一个节点,则返回前置节点,否则返回null
            if (index > 0) {
                return sortedChildrenNodes.get(index - 1);
            }

            return null;
        } catch (KeeperException | InterruptedException e) {
            throw new IllegalMonitorStateException(String.format("root path %s is not exists", ROOT_PATH));
        }
    }
}
